package com.cn.algorithm.one.tree;

import java.util.Comparator;

public class AVL<E> extends BBST<E> {

    public AVL(Comparator<E> comparator) {
        super(comparator);
    }

    public AVL() {
        this(null);
    }

    @Override
    public void afterPut(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                updateHeight(node);
            } else {
                reBalance(node);
                return;
            }
        }
    }

    public boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    public void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    public void reBalance(Node<E> grand) {
        Node<E> parent = ((AVLNode<E>) grand).tallerNode();
        Node<E> child = ((AVLNode<E>) parent).tallerNode();
        if (parent.isLeftChild()) {
            if (child.isRightChild()) {
                rotateLeft(parent);
            }
            rotateRight(grand);
        } else {
            if (child.isLeftChild()) {
                rotateRight(parent);
            }
            rotateLeft(grand);
        }
    }


    @Override
    public Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    @Override
    public void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);
        ((AVLNode<E>) grand).updateHeight();
        ((AVLNode<E>) parent).updateHeight();
    }

    public static class AVLNode<E> extends Node<E> {

        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        public int balanceFactor() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = Math.max(leftHeight, rightHeight) + 1;
        }

        public Node<E> tallerNode() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            if (leftHeight > rightHeight) {
                return left;
            } else if (leftHeight < rightHeight) {
                return right;
            } else {
                if (this.isLeftChild()) {
                    return right;
                } else {
                    return left;
                }
            }
        }
    }
}
